import os

import numpy as np

import caffe
import lmdb

def test_lmdb():
    lmdb_path = '/home/sunzy/workspace/data/MedlatTrainingData/lmdb'
    img_lmdb = lmdb.open(os.path.join(lmdb_path, "subints_train_lmdb"))
    txn = img_lmdb.begin()
    cursor = txn.cursor()
    datum = caffe.proto.caffe_pb2.Datum()
    num = 0
    num2 = 0
    for (idx, (key, value)) in enumerate(cursor):
        datum.ParseFromString(value)
        flat_x = np.fromstring(datum.data, dtype=np.float64)
        # print(key)
        x = flat_x.reshape(datum.channels, datum.height, datum.width)
        y = datum.label
        num += 1

        print(idx, key, x.shape, y)
        if y == 1:
            num2 += 1
    print(num, num2)

def main():
    test_lmdb()

if __name__ == "__main__":
    main()